import torch.nn as nn


class CardioMLP(nn.Module):
    def __init__(self):
        super(CardioMLP, self).__init__()
        self.layer1 = nn.Linear(10, 128)
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(0.2)

        self.layer2 = nn.Linear(128, 64)
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(0.2)

        self.layer3 = nn.Linear(64, 32)
        self.relu3 = nn.ReLU()
        self.dropout3 = nn.Dropout(0.6)

        self.output_layer = nn.Linear(32, 2)

    def forward(self, x):
        x = self.dropout1(self.relu1(self.layer1(x)))
        x = self.relu2(self.layer2(x))
        x = self.relu3(self.layer3(x))
        x = self.output_layer(x)
        return x


